# core/extractor.py
from typing import Set, Dict, Optional, List, Tuple
import sqlparse
from sqlparse.sql import (
    Identifier, IdentifierList, TokenList, Function,
    Parenthesis, Token, Case
)
from sqlparse.tokens import Keyword, Name, DML, Whitespace, Punctuation

from .models import StatementLineage, LineageResult
from .parser import parse_sql, statement_type

READ_KWS: Set[str]  = {"FROM"}
WRITE_KWS: Set[str] = {"INSERT", "UPDATE", "CREATE", "REPLACE"}
STOP_KWS: Set[str]  = {"WHERE", "GROUP", "ORDER", "LIMIT", "HAVING", "QUALIFY", "WINDOW",
                       "ON", "UNION", "INTERSECT", "EXCEPT"}


class LineageExtractor:
    def extract(self, sql: str) -> LineageResult:
        result = LineageResult()
        for st in parse_sql(sql):
            res = self._extract_statement(st)
            result.statements.append(res)
        return result

    # ---------- helpers ----------
    @staticmethod
    def _iter_identifiers(token: TokenList):
        """递归产出所有 Identifier 节点。"""
        if isinstance(token, Identifier):
            yield token
        elif isinstance(token, IdentifierList):
            for t in token.get_identifiers():
                if isinstance(t, Identifier):
                    yield t
                else:
                    yield from LineageExtractor._iter_identifiers(t)
        elif isinstance(token, (Function, Parenthesis, Case, TokenList)):
            for t in getattr(token, "tokens", []):
                yield from LineageExtractor._iter_identifiers(t)

    @staticmethod
    def _identifier_is_subquery(idf: Identifier) -> bool:
        """Identifier 内是否包含子查询括号。"""
        return any(isinstance(t, Parenthesis) for t in getattr(idf, "tokens", []))

    @staticmethod
    def _identifier_to_table_name(idf: Identifier) -> Optional[str]:
        """
        规范化为 catalog.schema.table / schema.table / table。
        逐 token 拼接 Name；遇到空白/AS/括号/非点标点停止（后面通常是别名）。
        对 (subquery) alias 返回 None。
        """
        if LineageExtractor._identifier_is_subquery(idf):
            return None

        parts: List[str] = []
        seen_name = False
        for t in idf.tokens:
            if seen_name and (t.ttype is Whitespace or (t.ttype is Keyword and t.value.upper() == "AS")):
                break
            if t.ttype is Name:
                parts.append(t.value)
                seen_name = True
                continue
            if t.ttype is Punctuation and t.value == ".":
                continue
            if seen_name:
                if isinstance(t, Parenthesis):
                    break
                if t.ttype is Punctuation and t.value != ".":
                    break

        if not parts:
            return None
        return ".".join(parts)

    @staticmethod
    def _token_is_kw(tk: Token, text: str) -> bool:
        return (tk.ttype is Keyword) and (tk.value.upper() == text)

    @staticmethod
    def _token_is_any_kw(tk: Token, cand: Set[str]) -> bool:
        return (tk.ttype is Keyword) and (tk.value.upper() in cand)

    # ---------- core ----------
    def _extract_statement(self, st) -> StatementLineage:
        op = statement_type(st)
        s = StatementLineage(op=op, raw_sql=str(st).strip())

        # 目标表 / 来源表
        self._collect_tables(st, s, mode="write", record_alias=True)
        self._collect_tables(st, s, mode="read",  record_alias=True)

        # 字段 & 列到列血缘
        if op == "SELECT" or "SELECT" in str(st).upper():
            self._collect_columns_and_lineage(st, s)

        s.tables |= (s.source_tables | s.target_tables)
        return s

    def _collect_tables(self, st: TokenList, s: StatementLineage, mode: str, record_alias: bool = True):
        tokens = list(st.tokens)
        n = len(tokens)
        i = 0
        while i < n:
            tk = tokens[i]

            # -------- 读取来源表：FROM / 任意 *JOIN --------
            if mode == "read" and (
                self._token_is_any_kw(tk, READ_KWS) or
                (tk.ttype is Keyword and "JOIN" in tk.value.upper())
            ):
                j = i + 1
                while j < n and (tokens[j].ttype in (Whitespace, Punctuation)):
                    j += 1
                if j >= n:
                    i += 1
                    continue
                seg = tokens[j]

                # (子查询)：递归解析内部，并记录子查询别名 -> 内部来源表
                if isinstance(seg, Parenthesis):
                    before = set(s.source_tables)
                    self._recurse_parenthesis_for_tables(seg, s, mode="read")

                    alias_name = None
                    k = j + 1
                    while k < n and (tokens[k].ttype in (Whitespace, Punctuation)):
                        k += 1

                    alias_name = None
                    if k < n:
                        if isinstance(tokens[k], Identifier):
                            alias_name = tokens[k].get_name() or tokens[k].get_alias() or tokens[k].value
                        elif tokens[k].ttype is Name:
                            alias_name = tokens[k].value  # ← 支持 Name 作为括号后别名

                    added = set(s.source_tables) - before
                    if alias_name and added:
                        subq_map = s.ctes.setdefault("_subq", {})
                        subq_map[alias_name] = sorted(added)

                    i = k + 1 if alias_name else j + 1
                    continue

                # Identifier / IdentifierList
                for idf in self._iter_identifiers(seg):
                    if self._identifier_is_subquery(idf):
                        for t in idf.tokens:
                            if isinstance(t, Parenthesis):
                                before = set(s.source_tables)
                                self._recurse_parenthesis_for_tables(t, s, mode="read")
                                break
                        continue

                    tbl = self._identifier_to_table_name(idf)
                    if tbl:
                        s.source_tables.add(tbl)
                        alias = idf.get_alias()
                        if record_alias and alias:
                            s.ctes.setdefault("_alias", set()).add(f"{alias}::{tbl}")

                i = j + 1
                continue

            # -------- 写入目标表：INSERT / UPDATE / CREATE / REPLACE --------
            if mode == "write" and self._token_is_any_kw(tk, WRITE_KWS):
                upper = tk.value.upper()
                j = i + 1

                if upper == "INSERT":
                    while j < n and not (isinstance(tokens[j], Identifier) or isinstance(tokens[j], IdentifierList)):
                        if self._token_is_kw(tokens[j], "INTO"):
                            j += 1
                            break
                        j += 1
                    while j < n and (tokens[j].ttype in (Whitespace, Punctuation)):
                        j += 1
                    if j < n:
                        for idf in self._iter_identifiers(tokens[j]):
                            tbl = self._identifier_to_table_name(idf)
                            if tbl:
                                s.target_tables.add(tbl)

                elif upper == "UPDATE":
                    while j < n and (tokens[j].ttype in (Whitespace, Punctuation)):
                        j += 1
                    if j < n:
                        for idf in self._iter_identifiers(tokens[j]):
                            tbl = self._identifier_to_table_name(idf)
                            if tbl:
                                s.target_tables.add(tbl)

                elif upper in {"CREATE", "REPLACE"}:
                    while j < n and not self._token_is_kw(tokens[j], "TABLE"):
                        j += 1
                    if j < n:
                        j += 1
                        while j < n and (tokens[j].ttype in (Whitespace, Punctuation)):
                            j += 1
                        if j < n:
                            for idf in self._iter_identifiers(tokens[j]):
                                tbl = self._identifier_to_table_name(idf)
                                if tbl:
                                    s.target_tables.add(tbl)

                i += 1
                continue

            i += 1

    def _recurse_parenthesis_for_tables(self, paren: Parenthesis, s: StatementLineage, mode: str):
        """在 ( ... ) 中如果含有 SELECT，则递归抽取 FROM/JOIN 段的表（不外泄内层 alias）。"""
        val = paren.value
        if "SELECT" not in val.upper():
            return
        inner = val[1:-1]
        for st in sqlparse.parse(inner):
            self._collect_tables(st, s, mode=mode, record_alias=False)

    # ---------- columns + column-to-column lineage ----------
    def _collect_columns_and_lineage(self, st: TokenList, s: StatementLineage):
        # 顶层别名映射（b -> 表）
        alias_map: Dict[str, str] = {}
        for item in s.ctes.get("_alias", set()):
            alias, full = item.split("::", 1)
            alias_map[alias] = full

        subq_map: Dict[str, List[str]] = s.ctes.get("_subq", {}) if isinstance(s.ctes, dict) else {}

        # 顶层 SELECT 段
        select_segment: List[Token] = []
        recording = False
        for tk in st.tokens:
            if tk.ttype is DML and tk.value.upper() == "SELECT":
                recording = True
                continue
            if recording:
                if tk.ttype is Keyword and tk.value.upper() == "FROM":
                    break
                select_segment.append(tk)
        if not select_segment:
            return

        segment = TokenList(select_segment)

        FUNC_BLACKLIST = {
            "count","sum","avg","min","max","date","cast","substring","substr",
            "date_add","date_sub","date_format","from_unixtime","current_date",
            "coalesce","nvl","if","case","when","then","else","end"
        }

        def _add_tbl_col(tbl_alias: Optional[str], col_name: Optional[str]):
            """表→列（columns_by_table）"""
            if not col_name or col_name == "*":
                return
            if tbl_alias is None and col_name.lower() in FUNC_BLACKLIST:
                return
            if tbl_alias:
                tbl_full = alias_map.get(tbl_alias) or tbl_alias
                s.columns_by_table.setdefault(tbl_full, set()).add(col_name)
            else:
                s.columns_by_table.setdefault("*", set()).add(col_name)

        def _gather_sources_from_node(node: TokenList) -> List[Tuple[str,str]]:
            """
            收集一个投影表达式中的源列，返回 [(表全名/或subquery:别名, 列名)]。
            优先使用 Identifier，兜底 Name '.' Name。
            """
            pairs: List[Tuple[str,str]] = []

            # A) Identifier 方式
            for idf in self._iter_identifiers(node):
                alias = idf.get_parent_name()
                col   = idf.get_real_name() or idf.get_name()
                if not col or col == "*":
                    continue
                if alias in alias_map:
                    pairs.append((alias_map[alias], col))
                elif alias in subq_map:
                    # 源来自子查询别名：标记成 subquery:alias
                    pairs.append((f"subquery:{alias}", col))
                else:
                    # 无法解析的别名/裸列：忽略（避免误报）
                    pass

            # B) Name '.' Name 兜底
            def _walk(tok: TokenList):
                toks = getattr(tok, "tokens", []) or []
                i = 0
                while i+2 < len(toks):
                    t0, t1, t2 = toks[i], toks[i+1], toks[i+2]
                    if t0.ttype is Name and (t1.ttype is Punctuation and t1.value == ".") and t2.ttype is Name:
                        a, c = t0.value, t2.value
                        if a in alias_map:
                            pairs.append((alias_map[a], c))
                        elif a in subq_map:
                            pairs.append((f"subquery:{a}", c))
                        i += 3
                        continue
                    # 递归
                    if isinstance(t0, (TokenList, Parenthesis, Identifier, IdentifierList, Function, Case)):
                        _walk(t0)
                    i += 1
                # 收尾节点递归
                if toks:
                    last = toks[-1]
                    if isinstance(last, (TokenList, Parenthesis, Identifier, IdentifierList, Function, Case)):
                        _walk(last)

            _walk(node)
            return pairs

        # 1) 表→列：先把所有列统计好（与之前行为一致）
        for node in segment.tokens:
            for idf in self._iter_identifiers(node):
                _add_tbl_col(idf.get_parent_name(), idf.get_real_name() or idf.get_name())
        # 再做兜底
        def _walk_for_tblcols(tok: TokenList):
            toks = getattr(tok, "tokens", []) or []
            i = 0
            while i+2 < len(toks):
                t0, t1, t2 = toks[i], toks[i+1], toks[i+2]
                if t0.ttype is Name and (t1.ttype is Punctuation and t1.value == ".") and t2.ttype is Name:
                    _add_tbl_col(t0.value, t2.value)
                    i += 3
                    continue
                if isinstance(t0, (TokenList, Parenthesis, Identifier, IdentifierList, Function, Case)):
                    _walk_for_tblcols(t0)
                i += 1
            if toks:
                last = toks[-1]
                if isinstance(last, (TokenList, Parenthesis, Identifier, IdentifierList, Function, Case)):
                    _walk_for_tblcols(last)
        _walk_for_tblcols(segment)

        # 2) 列→列：为每个“投影项”提取输出名（alias 或 name），并链接源列
        # 把 SELECT 段拆成每个“投影项”
        items: List[TokenList] = []
        # 如果存在 IdentifierList，优先用它的 children；否则退化为逐 token 粗拆
        idlist = None
        for t in segment.tokens:
            if isinstance(t, IdentifierList):
                idlist = t
                break
        if idlist:
            for it in idlist.get_identifiers():
                items.append(it)  # Identifier / Function / 等
        else:
            # 粗拆：逗号分段
            cur = []
            for t in segment.tokens:
                if t.ttype is Punctuation and t.value == ",":
                    if cur:
                        items.append(TokenList(cur))
                        cur = []
                else:
                    cur.append(t)
            if cur:
                items.append(TokenList(cur))

        for item in items:
            # 输出名（alias）
            out_alias = None
            if isinstance(item, Identifier):
                out_alias = item.get_alias() or item.get_name()
            else:
                # 处理 "表达式 AS 别名"：在 item 内部寻找 AS/Name
                toks = getattr(item, "tokens", []) or []
                for i, tk in enumerate(toks):
                    if tk.ttype is Keyword and tk.value.upper() == "AS":
                        # AS 后第一个 Name/Identifier 的名字
                        j = i + 1
                        while j < len(toks) and toks[j].ttype in (Whitespace, Punctuation):
                            j += 1
                        if j < len(toks):
                            if isinstance(toks[j], Identifier):
                                out_alias = toks[j].get_name() or toks[j].get_alias() or toks[j].value
                            elif toks[j].ttype is Name:
                                out_alias = toks[j].value
                        break
                # 如无 AS，尝试末尾 Identifier 作为别名
                if not out_alias and toks:
                    last = toks[-1]
                    if isinstance(last, Identifier):
                        out_alias = last.get_name() or last.get_alias()

            if not out_alias:
                # 没有别名就不连“列到列”，以免把同名直通列造成噪声
                continue

            # 源列集合
            src_pairs = _gather_sources_from_node(item)  # [(full_table, col)]
            for full_tbl, col in src_pairs:
                key = f"{full_tbl}.{col}"
                s.column_lineage.setdefault(key, set()).add(out_alias)
